#Matan Keidar 066500992
#Eliahu Khalastchi 040854705

from Corpus import Corpus

class BigramCorpus():
    '''
    Represents a corpus that holds bigrams of words  
    '''


    def __init__(self, fileName, uniCorpus):
        '''
        Constructor
        '''
        self.unigramCorpus = uniCorpus
        self.bigramDic = {}
        self.s = 0
        self.x = 0
        
        #check if input file was given as argument
        if fileName is None:
            return
        
        #open input file
        input = open(fileName, "r");
        
        # read first word in file
        word1 = input.readline().strip()    # w (the)
        
        #add each bigram words in input to the corpus
        for line in input:  # (the boy)
            word2 = line.strip()    # w' (boy)

            self.addWords(word2, word1) # w',w (boy, the) 
            word1 = word2
        
#        self.addWordUnigram(word1)        
        input.close()
        
    def S(self):
        '''
        Gets the total size of words that were added to the corpus
        '''
        
        if self.s == 0:
            self.s=sum(self.bigramDic.values())
        return self.s

    def X(self):
        '''
        Gets the size of unique words that were added to the corpus
        '''
        
        if self.x == 0:
            self.x=len(self.bigramDic)
        return self.x
        
    def getWordByIndex(self, i):
        return self.allWords.keys()[i]
    
#    def addWordUnigram(self, word):
#        index = self.allWords[word]
#        self.unigramCorpus.addWord(index)
        
    def addWords(self, word2, word1):
        '''
        adds the given (w2,w1) to corpus 
        w1 - condition word
        w2 - word that is conditioned by w1 
        '''
        
        if self.bigramDic.get((word2, word1)) is None:
            self.bigramDic[(word2, word1)] = 0
        
        self.bigramDic[(word2, word1)] += 1
        
    def C(self,words):
        '''
        Gets the frequency of the given (w',w) 
        '''
        c = self.bigramDic.get(words)
        if c is None:
            return 0
        return c
    
    def getEntries(self):
        '''
        Gets a collection of words that are stored in the corpus
        '''        
        return self.bigramDic.keys()
    